# -*- coding:utf-8 -*-
# Author:lixuecheng
import paramiko
import re
from pro.core.logger import log, logger
import os
import time




class DoSSH:
    def __init__(self, host, port, user, passwd=None, key=None, key_passwd=None):
        transport = paramiko.Transport((str(host), int(port)))
        if key is None and passwd is None:
            raise Exception('passwd or key must input one')
        elif key is None:
            transport.connect(username=str(user), password=str(passwd))
        else:
            if key_passwd is None:
                private_key = paramiko.RSAKey.from_private_key_file(str(key))
            else:
                private_key = paramiko.RSAKey.from_private_key_file(
                    str(key), str(key_passwd))
            transport.connect(username=user, pkey=private_key)
        ssh = paramiko.SSHClient()
        # 允许连接不在know_hosts文件中的主机
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh._transport = transport
        sftp = ssh.open_sftp()
        self.ssh_value = ssh
        self.sftp_value = sftp
        self.transport = transport

    def ssh(self, cmd, inp=None):
        stdin, stdout, stderr = self.ssh_value.exec_command(str(cmd))
        if inp is not None:
            stdin.write(str(inp))
        if len(stderr.read().decode()) > 0:
            raise Exception(stderr.read().decode())
        return stdout.read().decode()

    def ssh_console(self):
        shell = self.ssh_value.invoke_shell()
        shell.settimeout(1)

        command=''
        while True:
            try:
                recv = shell.recv(512).decode()

                if recv:
                    
                        # print(recv,222)
                    rec = re.sub(b'\x1b'.decode()+r'.+?m', '', recv)
                    # aa=re.findall(r'\?.+?m',recv)
                   
                    if rec.startswith(command.replace('\n','')):
                        rec=rec.replace(command.replace('\n',''),'')
                    print(rec, end='')
                        # print('_'*30)
                        # print(aa)

                else:
                    print(recv, 112)
                    continue
            except Exception as e:
                if e is KeyboardInterrupt:
                    break
                try:
                    command = input() + "\n"
                except KeyboardInterrupt:
                    shell.send('exit\n')
                    recv = shell.recv(512).decode()
                    if recv:
                        rec = re.compile(r'\?\[.+?m').sub('', recv)

                        print(rec, end='')
                    else:
                        continue
                    break

                shell.send(command)


                if command == 'exit\n':
                    recv = shell.recv(512).decode()
                    if recv:
                        rec = re.compile(r'\?\[.+?m').sub('', recv)

                        print(rec, end='')
                    else:
                        continue
                    break

    def sftp(self):
        return self.sftp_value

    def put(self, local_file, remote_file):
        if not os.path.exists(local_file):
            raise Exception(local_file + ',文件不存在')
        if not self.is_exist(os.path.dirname(remote_file)):
            self.mkdirs(os.path.dirname(remote_file))
        self.sftp_value.put(local_file, remote_file)

    def get(self, local_file, remote_file):
        if os.path.exists(local_file):
            t = time.strftime('%Y%m%d_%H%M%S', time.localtime())
            os.rename(local_file, local_file + '_backup_' + t)
        else:
            if not os.path.exists(os.path.dirname(local_file)):
                raise Exception(os.path.dirname(local_file) + ',目录不存在')
        if not self.is_exist(remote_file):
            raise Exception(remote_file + '文件不存在')
        self.sftp_value.get(remote_file, local_file)

    def is_exist(self, path):
        try:
            self.ssh('ls ' + str(path))
            return True
        except:
            return False

    def mkdirs(self, path):

        if self.is_exist(path):
            pass
        else:
            if path == '/':
                raise Exception('创建文件夹失败')
            else:
                try:
                    self.sftp_value.mkdir(path)
                    self.log.add('创建文件夹成功：' + path)
                    print('创建文件夹成功：', path)
                except:
                    self.mkdirs(os.path.dirname(path))
                    self.mkdirs(path)

    def close(self):
        self.sftp_value.close()
        self.ssh_value.close()
        self.transport.close()
        
# ssh.connect(
#     hostname="172.16.32.210",
#     port=22,
#     username="root",
#     password="Test123456!"
# )
# d=DoSSH('172.16.32.210',22,'root',"Test123456!")
# # d.ssh_console()
# a=d.ssh('ls -l')
# print(a)
# d.close()
